﻿// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

#pragma warning disable SA1512 // Single-line comments should not be followed by blank line
#pragma warning disable SA1604 // Element documentation should have summary
// define TRACE_LEAKS to get additional diagnostics that can lead to the leak sources. note: it will
// make everything about 2-3x slower
//
// #define TRACE_LEAKS

// define DETECT_LEAKS to detect possible leaks
// #if DEBUG
// #define DETECT_LEAKS  //for now always enable DETECT_LEAKS in debug.
// #endif

namespace Microsoft.Testing.Platform.Helpers;

/// <summary>
/// Generic implementation of object pooling pattern with predefined pool size limit. The main
/// purpose is that limited number of frequently used objects can be kept in the pool for
/// further recycling.
///
/// Notes:
/// 1) it is not the goal to keep all returned objects. Pool is not meant for storage. If there
///    is no space in the pool, extra returned objects will be dropped.
///
/// 2) it is implied that if object was obtained from a pool, the caller will return it back in
///    a relatively short time. Keeping checked out objects for long durations is ok, but
///    reduces usefulness of pooling. Just new up your own.
///
/// Not returning objects to the pool in not detrimental to the pool's work, but is a bad practice.
/// Rationale:
///    If there is no intent for reusing the object, do not use pool - just use "new".
/// </summary>
#pragma warning disable SA1214 // Readonly fields should appear before non-readonly fields
[ExcludeFromCodeCoverage]
#pragma warning disable SA1618 // Generic type parameters should be documented

#pragma warning disable SA1201 // Elements should appear in the correct order
internal sealed class ObjectPool<T>
#pragma warning restore SA1618 // Generic type parameters should be documented
    where T : class
{
    [DebuggerDisplay("{Value,nq}")]
#pragma warning disable CA1815 // Override equals and operator equals on value types
#pragma warning disable CA1051 // Do not declare visible instance fields
    private struct Element
    {
#pragma warning disable IDE1006 // Naming Styles
        internal T? Value;
#pragma warning restore IDE1006 // Naming Styles
    }
#pragma warning restore CA1051 // Do not declare visible instance fields
#pragma warning restore CA1815 // Override equals and operator equals on value types

    /// <remarks>
    /// Not using System.Func{T} because this file is linked into the (debugger) Formatter,
    /// which does not have that type (since it compiles against .NET 2.0).
    /// </remarks>
    internal delegate T Factory();

    // Storage for the pool objects. The first item is stored in a dedicated field because we
    // expect to be able to satisfy most requests from it.
    private T? _firstItem;
    private readonly Element[] _items;

    // factory is stored for the lifetime of the pool. We will call this only when pool needs to
    // expand. compared to "new T()", Func gives more flexibility to implementers and faster
    // than "new T()".
    private readonly Factory _factory;

#if DETECT_LEAKS
    private static readonly ConditionalWeakTable<T, LeakTracker> leakTrackers = new ConditionalWeakTable<T, LeakTracker>();

    private class LeakTracker : IDisposable
    {
        private volatile bool disposed;

#if TRACE_LEAKS
        internal volatile object Trace = null;
#endif

        public void Dispose()
        {
            disposed = true;
            GC.SuppressFinalize(this);
        }

        private string GetTrace()
        {
#if TRACE_LEAKS
            return Trace == null ? "" : Trace.ToString();
#else
            return "Leak tracing information is disabled. Define TRACE_LEAKS on ObjectPool`1.cs to get more info \n";
#endif
        }

        ~LeakTracker()
        {
            if (!this.disposed && !Environment.HasShutdownStarted)
            {
                var trace = GetTrace();

                // If you are seeing this message it means that object has been allocated from the pool
                // and has not been returned back. This is not critical, but turns pool into rather
                // inefficient kind of "new".
                Debug.WriteLine($"TRACEOBJECTPOOLLEAKS_BEGIN\nPool detected potential leaking of {typeof(T)}. \n Location of the leak: \n {GetTrace()} TRACEOBJECTPOOLLEAKS_END");
            }
        }
    }
#endif

    internal ObjectPool(Factory factory)
#pragma warning disable RS1035 // Do not use APIs banned for analyzers
#pragma warning disable RS0030 // Do not use banned APIs
        : this(factory, Environment.ProcessorCount * 2)
#pragma warning restore RS0030 // Do not use banned APIs
#pragma warning restore RS1035 // Do not use APIs banned for analyzers
    {
    }

    internal ObjectPool(Factory factory, int size)
    {
        RoslynDebug.Assert(size >= 1);
        _factory = factory;
        _items = new Element[size - 1];
    }

    private T CreateInstance()
    {
        T inst = _factory();
        return inst;
    }

    /// <summary>
    /// Produces an instance.
    /// </summary>
    /// <remarks>
    /// Search strategy is a simple linear probing which is chosen for it cache-friendliness.
    /// Note that Free will try to store recycled objects close to the start thus statistically
    /// reducing how far we will typically search.
    /// </remarks>
    internal T Allocate()
    {
        // PERF: Examine the first element. If that fails, AllocateSlow will look at the remaining elements.
        // Note that the initial read is optimistically not synchronized. That is intentional.
        // We will interlock only when we have a candidate. in a worst case we may miss some
        // recently returned objects. Not a big deal.
        T? inst = _firstItem;
        if (inst == null || inst != Interlocked.CompareExchange(ref _firstItem, null, inst))
        {
            inst = AllocateSlow();
        }

#if DETECT_LEAKS
        var tracker = new LeakTracker();
        leakTrackers.Add(inst, tracker);

#if TRACE_LEAKS
        var frame = CaptureStackTrace();
        tracker.Trace = frame;
#endif
#endif
        return inst;
    }

    private T AllocateSlow()
    {
        Element[] items = _items;

        for (int i = 0; i < items.Length; i++)
        {
            // Note that the initial read is optimistically not synchronized. That is intentional.
            // We will interlock only when we have a candidate. in a worst case we may miss some
            // recently returned objects. Not a big deal.
            T? inst = items[i].Value;
            if (inst != null &&
                inst == Interlocked.CompareExchange(ref items[i].Value, null, inst))
            {
                return inst;
            }
        }

        return CreateInstance();
    }

    /// <summary>
    /// Returns objects to the pool.
    /// </summary>
    /// <remarks>
    /// Search strategy is a simple linear probing which is chosen for it cache-friendliness.
    /// Note that Free will try to store recycled objects close to the start thus statistically
    /// reducing how far we will typically search in Allocate.
    /// </remarks>
    internal void Free(T obj, CancellationToken cancellationToken)
    {
        // Do not free in presence of cancellation.
        // See https://github.com/dotnet/roslyn/issues/46859 for details.
        if (cancellationToken.IsCancellationRequested)
        {
            return;
        }

        Validate(obj);
        ForgetTrackedObject(obj);

        if (_firstItem == null)
        {
            // Intentionally not using interlocked here.
            // In a worst case scenario two objects may be stored into same slot.
            // It is very unlikely to happen and will only mean that one of the objects will get collected.
            _firstItem = obj;
        }
        else
        {
            FreeSlow(obj);
        }
    }

    private void FreeSlow(T obj)
    {
        Element[] items = _items;
        for (int i = 0; i < items.Length; i++)
        {
            if (items[i].Value == null)
            {
                // Intentionally not using interlocked here.
                // In a worst case scenario two objects may be stored into same slot.
                // It is very unlikely to happen and will only mean that one of the objects will get collected.
                items[i].Value = obj;
                break;
            }
        }
    }

    /// <summary>
    /// Removes an object from leak tracking.
    ///
    /// This is called when an object is returned to the pool.  It may also be explicitly
    /// called if an object allocated from the pool is intentionally not being returned
    /// to the pool.  This can be of use with pooled arrays if the consumer wants to
    /// return a larger array to the pool than was originally allocated.
    /// </summary>
    [Conditional("DEBUG")]
    internal static void ForgetTrackedObject(T old, T? replacement = null)
    {
#if DETECT_LEAKS
        LeakTracker tracker;
        if (leakTrackers.TryGetValue(old, out tracker))
        {
            tracker.Dispose();
            leakTrackers.Remove(old);
        }
        else
        {
            var trace = CaptureStackTrace();
            Debug.WriteLine($"TRACEOBJECTPOOLLEAKS_BEGIN\nObject of type {typeof(T)} was freed, but was not from pool. \n Callstack: \n {trace} TRACEOBJECTPOOLLEAKS_END");
        }

        if (replacement != null)
        {
            tracker = new LeakTracker();
            leakTrackers.Add(replacement, tracker);
        }
#endif
    }

#if DETECT_LEAKS
    private static Lazy<Type> _stackTraceType = new Lazy<Type>(() => Type.GetType("System.Diagnostics.StackTrace"));

    private static object CaptureStackTrace()
    {
        return Activator.CreateInstance(_stackTraceType.Value);
    }
#endif

    [Conditional("DEBUG")]
    private void Validate(object obj)
    {
        RoslynDebug.Assert(_firstItem != obj, "freeing twice?");

        Element[] items = _items;
        for (int i = 0; i < items.Length; i++)
        {
            T? value = items[i].Value;
            if (value == null)
            {
                return;
            }

            RoslynDebug.Assert(value != obj, "freeing twice?");
        }
    }
}
